package com.bw.test3;

import com.alibaba.alink.operator.batch.evaluation.EvalBinaryClassBatchOp;
import com.alibaba.alink.operator.batch.evaluation.EvalMultiClassBatchOp;
import com.alibaba.alink.operator.common.evaluation.BinaryClassMetrics;
import com.alibaba.alink.operator.common.evaluation.MultiClassMetrics;
import org.apache.flink.types.Row;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.classification.KnnPredictBatchOp;
import com.alibaba.alink.operator.batch.classification.KnnTrainBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class KnnTrainBatchOpTest {
	@Test
	public void testKnnTrainBatchOp() throws Exception {

		BatchOperator.setParallelism(1);
		// 数据集
		List <Row> df = Arrays.asList(
			Row.of(1, 0, 0),
			Row.of(2, 8, 8),
			Row.of(1, 1, 2),
			Row.of(2, 9, 10),
			Row.of(1, 3, 1),
			Row.of(2, 10, 7)
		);
		BatchOperator <?> dataSourceOp = new MemSourceBatchOp(df, "label int, f0 int, f1 int");

		BatchOperator <?> trainOp = new KnnTrainBatchOp().setFeatureCols("f0", "f1").setLabelCol("label")
			.setDistanceType("EUCLIDEAN").linkFrom(dataSourceOp);

		// 预测
		BatchOperator <?> predictOp = new KnnPredictBatchOp().setPredictionCol("pred").setPredictionDetailCol("pred_detail").setK(4).linkFrom(trainOp,
			dataSourceOp);

		predictOp.print();



		// 二分类评估
		BinaryClassMetrics metrics = new EvalBinaryClassBatchOp().setLabelCol("label").setPredictionDetailCol(
				"pred_detail").linkFrom(predictOp).collectMetrics();
		System.out.println("AUC:" + metrics.getAuc());
		System.out.println("KS:" + metrics.getKs());
		System.out.println("PRC:" + metrics.getPrc());
		System.out.println("Accuracy:" + metrics.getAccuracy());
		System.out.println("Macro Precision:" + metrics.getMacroPrecision());
		System.out.println("Micro Recall:" + metrics.getMicroRecall());
		System.out.println("Weighted Sensitivity:" + metrics.getWeightedSensitivity());


		// 多分类
		MultiClassMetrics metrics1 = new EvalMultiClassBatchOp().setLabelCol("label").setPredictionDetailCol(
				"pred_detail").linkFrom(predictOp).collectMetrics();

		System.out.println("Prefix0 accuracy:" + metrics1.getAccuracy());
		System.out.println("Prefix1 recall:" + metrics1.getRecall("1"));// todo 查一下1的意思
		System.out.println("Macro Precision:" + metrics1.getMacroPrecision());
		System.out.println("Micro Recall:" + metrics1.getMicroRecall());
		System.out.println("Weighted Sensitivity:" + metrics1.getWeightedSensitivity());

	}
}